import sys
from Prediction import *
from utils.DorisQueryTool import DorisQueryTool
from utils.MysqlQueryTool import MysqlQueryTool
import pandas as pd
import os
import matplotlib.pyplot as plt


def plot_predictions(true_dates, true_values, pred_values, title="predict_flow_compare"):
    """
    绘制真实值与预测值的对比折线图
    :param true_dates: 日期列表 (datetime格式)
    :param true_values: 真实值列表
    :param pred_values: 预测值列表
    :param title: 图表标题
    """
    plt.figure(figsize=(25, 10))

    # 绘制曲线
    plt.plot(true_dates, true_values, 'b-', label='flow', marker='o', linewidth=1, markersize=1)
    plt.plot(true_dates, pred_values, 'r--', label='predict_flow', marker='s', linewidth=1,markersize=1)

    # 添加标注
    # for i, (true, pred) in enumerate(zip(true_values, pred_values)):
    #     plt.text(true_dates[i], max(true, pred), f"{pred:.0f}",
    #              ha='center', va='bottom', fontsize=9)

    # 美化图表
    plt.title(title, fontsize=14)
    plt.xlabel('date', fontsize=12)
    plt.ylabel('flow', fontsize=12)
    plt.xticks(rotation=90)
    plt.grid(linestyle='--', alpha=0.7)
    plt.legend()
    plt.tight_layout()

    # 保存图片
    plt.savefig('prediction_comparison.png', dpi=300)
    plt.show()


if __name__ == "__main__":
    # 查询所有收费站近两年数据
    print("Received arguments:", sys.argv)
    # 提取具体参数
    roadId = sys.argv[1] if len(sys.argv) > 1 else "G0004130010"
    stationId = sys.argv[2] if len(sys.argv) > 2 else "G0004130010150"
    intervalId = sys.argv[3] if len(sys.argv) > 3 else "G000413001000220"
    portalIds = sys.argv[4] if len(sys.argv) > 4 else "'G000413001000220010','G000413001000220020'"
    type = sys.argv[5] if len(sys.argv) > 5 else "0"
    #modelPath = sys.argv[6] if len(sys.argv) > 6 else "E:\\data\\G0001130010\\G0001130010030\\6a9e9087-09c9-47aa-b220-a58f62f30003.pth"
    modelPath = sys.argv[6] if len(sys.argv) > 6 else "E:\\data\\aaa.pth"
    startTime = sys.argv[7] if len(sys.argv) > 7 else "2025-01-01"
    endTime = sys.argv[8] if len(sys.argv) > 8 else "2025-06-20"

    train_ratio = sys.argv[9] if len(sys.argv) > 9 else "75"
    val_ratio = sys.argv[10] if len(sys.argv) > 10 else "15"
    test_ratio = sys.argv[11] if len(sys.argv) > 11 else "15"
    batch_size = sys.argv[12] if len(sys.argv) > 12 else "28"
    input_length = sys.argv[13] if len(sys.argv) > 13 else "28"
    output_length = sys.argv[14] if len(sys.argv) > 14 else "7"
    learning_rate = sys.argv[15] if len(sys.argv) > 15 else "0.001"
    num_blocks = sys.argv[16] if len(sys.argv) > 16 else "2"
    dim = sys.argv[17] if len(sys.argv) > 17 else "256"
    scalerPath = sys.argv[18] if len(sys.argv) > 18 else "E:\\data\\aaa.joblib"
    # modelPath = os.path.normpath(modelPath)
    # scalerPath = os.path.normpath(scalerPath)

    print("Received roadId:", roadId)
    print("Received stationId:", stationId)
    print("Received intervalId:", intervalId)
    print("Received portalIds:", portalIds)
    print("Received type:", type)
    print("Received modelPath:", modelPath)
    print("Received startTime:", startTime)
    print("Received endTime:", endTime)

    print("Received train_ratio:", train_ratio)
    print("Received val_ratio:", val_ratio)
    print("Received test_ratio:", test_ratio)
    print("Received batch_size:", batch_size)
    print("Received input_length:", input_length)
    print("Received output_length:", output_length)
    print("Received learning_rate:", learning_rate)
    print("Received num_blocks:", num_blocks)
    print("Received dim:", dim)
    type = int(type)
    param = {}
    param["train_ratio"] = int(train_ratio) / 100
    param["val_ratio"] = int(val_ratio) / 100
    param["test_ratio"] = int(test_ratio) / 100
    param["batch_size"] = int(batch_size)
    param["input_length"] = int(input_length)
    param["output_length"] = int(output_length)
    param["learning_rate"] = float(learning_rate)
    param["num_blocks"] = int(num_blocks)
    param["dim"] = int(dim)

    # 加载模型
    model = load_pretrained_model(modelPath,param)

    start_time = datetime.strptime(startTime, "%Y-%m-%d")
    end_time = datetime.strptime(endTime, "%Y-%m-%d")

    with MysqlQueryTool() as tool:
        # 示例1：返回DataFrame
        holiday = tool.execute_query("SELECT * FROM `holiday` where is_holiday = 1 order BY date  ")

    # 查询预测时间前 input_length 数据
    search_start_time = (start_time - timedelta(days=param["input_length"])).strftime("%Y-%m-%d %H:%M:%S")
    search_end_time = (end_time + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S")
    # 查询路况
    with MysqlQueryTool() as tool:
        # 示例1：返回DataFrame
        stationInfo = tool.execute_query(
            "SELECT status station_status,total_hour station_hour,road_id,station_id,date FROM `accident_station_info` where road_id = " + roadId + " and station_id = " + stationId + " and date >= '" + search_start_time + "' and date <= '" + search_end_time + "' order BY date  ")
        print(stationInfo)

    with DorisQueryTool() as tool:
        if type == 0:
            df = tool.execute_query(
                "SELECT count( 1 ) flow,EN_TOLL_SECTION_ID roadId,EN_TOLL_STATION_ID stationId,DATE_FORMAT( EN_TIME, '%Y-%m-%d' ) date FROM `toll_en` WHERE EN_TIME >= '" + search_start_time + "' AND EN_TIME < '" + search_end_time + "' AND EN_TOLL_SECTION_ID = '" + roadId + "' AND EN_TOLL_STATION_ID = '" + stationId + "' GROUP BY roadId,stationId,date ")
            group = df.sort_values('date')
            # 补充缺失数据
            group['date'] = pd.to_datetime(group['date'])  # 确保日期是datetime类型
            # 创建完整的日期范围
            date_range = pd.date_range(start=group['date'].min(), end=group['date'].max(), freq='D')
            # 重新索引DataFrame
            group = group.set_index('date').reindex(date_range, fill_value=0).rename_axis('date').reset_index()
            print(group)
        if type == 1:
            if roadId is None or stationId is None:
                df = tool.execute_query(
                    "SELECT count( 1 ) flow,DATE_FORMAT(EX_TIME, '%Y-%m-%d') date FROM `toll_ex` where EX_TIME >= '" + search_start_time + "' and EX_TIME < '" + search_end_time + "'   group by date")
            else:
                df = tool.execute_query(
                    "SELECT count( 1 ) flow,EX_TOLL_SECTION_ID roadId,EX_TOLL_STATION_ID stationId,DATE_FORMAT(EX_TIME, '%Y-%m-%d') date FROM `toll_ex` where EX_TIME >= '"+search_start_time+"' and EX_TIME < '"+search_end_time+"' and EX_TOLL_SECTION_ID = '"+roadId+"' and EX_TOLL_STATION_ID = '"+stationId+"'  group by EX_TOLL_SECTION_ID,EX_TOLL_STATION_ID ,date")
            group = df.sort_values('date')
            # 补充缺失数据
            group['date'] = pd.to_datetime(group['date'])  # 确保日期是datetime类型
            # 创建完整的日期范围
            date_range = pd.date_range(start=group['date'].min(), end=group['date'].max(), freq='D')
            # 重新索引DataFrame
            group = group.set_index('date').reindex(date_range, fill_value=0).rename_axis('date').reset_index()
            print(group)
        if type == 2 or type == 3:
            df = tool.execute_query(
                #"SELECT count( 1 ) flow,TOLL_SECTION_ID roadId,DATE_FORMAT( TRANS_TIME, '%Y-%m-%d' ) date FROM	`gantry_trade` WHERE TRANS_TIME >= '" + search_start_time + "' AND TRANS_TIME < '" + search_end_time + "' AND TOLL_SECTION_ID = '" + roadId + "' AND TOLL_GANTRY_ID IN ( " + portalIds + ") and c_master = 1 GROUP BY roadId,date")
                "SELECT count( 1 ) flow,TOLL_SECTION_ID roadId,DATE_FORMAT( TRANS_TIME, '%Y-%m-%d' ) date FROM	`gantry_trade` WHERE TRANS_TIME >= '" + search_start_time + "' AND TRANS_TIME < '" + search_end_time + "' AND TOLL_SECTION_ID = '" + roadId + "' AND TOLL_GANTRY_ID IN ( " + portalIds + ")  GROUP BY roadId,date")
            group = df.sort_values('date')
            # 补充缺失数据
            group['date'] = pd.to_datetime(group['date'])  # 确保日期是datetime类型
            # 创建完整的日期范围
            date_range = pd.date_range(start=group['date'].min(), end=group['date'].max(), freq='D')
            # 重新索引DataFrame
            group = group.set_index('date').reindex(date_range, fill_value=0).rename_axis('date').reset_index()
            print(group)
    # 预处理数据
    features_norm, scaler = preprocess_data(group,holiday,stationInfo,scalerPath)
    result = {}
    insert_datas = []
    while start_time <= end_time:
        prediction, dates = predict_future(model, scaler, start_time, group,param,holiday,scalerPath)
        # 7. 打印结果
        if (len(dates) > 0):
            dict = {"roadId": roadId, "stationId": stationId,"intervalId":intervalId,"type":type, "flowPrediction": int(prediction[0]),
                    "datePrediction": dates[0].strftime("%Y-%m-%d"),"datePredictionStr": dates[0].strftime("%Y-%m-%d")}
            # 查询真实车流量
            target_row = group[group['date'] == dates[0]]
            if not target_row.empty:
                dict["flow"] = int(target_row['flow'].values[0])
            else:
                dict["flow"] = 0
            insert_datas.append(dict)
        start_time = start_time + timedelta(days=1)
    # 打印预测结果
    print("flow prediction result:")
    print(insert_datas)
    datePredictionStrs = [item['datePredictionStr'] for item in insert_datas]
    flowPredictions = [item['flowPrediction'] for item in insert_datas]
    flows = [item['flow'] for item in insert_datas]
    plot_predictions(datePredictionStrs,flows,flowPredictions)



